import torch
from sklearn.metrics import accuracy_score, recall_score, f1_score

def train_model(model, train_loader, val_loader, num_epochs, criterion, optimizer, device, save_path='best_model.pth', scheduler=None, patience=5):
    best_val_score = 0.0
    no_improve_epochs = 0

    for epoch in range(num_epochs):
        model.train()
        train_loss, train_preds, train_labels = 0.0, [], []

        for time_data, space_data, freq_data, labels in train_loader:
            time_data, space_data, freq_data, labels = map(lambda x: x.to(device), [time_data, space_data, freq_data, labels])
            optimizer.zero_grad()
            outputs = model(time_data, space_data, freq_data)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * labels.size(0)
            train_preds.extend(outputs.argmax(dim=1).cpu().numpy())
            train_labels.extend(labels.cpu().numpy())

        train_loss /= len(train_loader.dataset)
        train_acc = accuracy_score(train_labels, train_preds)
        train_recall = recall_score(train_labels, train_preds, average='macro')
        train_f1 = f1_score(train_labels, train_preds, average='macro')

        # Validation
        model.eval()
        val_loss, val_preds, val_labels = 0.0, [], []
        with torch.no_grad():
            for time_data, space_data, freq_data, labels in val_loader:
                time_data, space_data, freq_data, labels = map(lambda x: x.to(device), [time_data, space_data, freq_data, labels])
                outputs = model(time_data, space_data, freq_data)
                loss = criterion(outputs, labels)
                val_loss += loss.item() * labels.size(0)
                val_preds.extend(outputs.argmax(dim=1).cpu().numpy())
                val_labels.extend(labels.cpu().numpy())

        val_loss /= len(val_loader.dataset)
        val_acc = accuracy_score(val_labels, val_preds)
        val_recall = recall_score(val_labels, val_preds, average='macro')
        val_f1 = f1_score(val_labels, val_preds, average='macro')

        print(f"Epoch {epoch+1}/{num_epochs} | Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} Recall: {train_recall:.4f} F1: {train_f1:.4f} | Val Loss: {val_loss:.4f} Acc: {val_acc:.4f} Recall: {val_recall:.4f} F1: {val_f1:.4f}")

        # Best model save (based on F1)
        if val_f1 > best_val_score:
            best_val_score = val_f1
            no_improve_epochs = 0
            torch.save(model.state_dict(), save_path)
            print(f"✅ Best model saved to {save_path}")
        else:
            no_improve_epochs += 1
            if no_improve_epochs >= patience:
                print("⏹️ Early stopping triggered.")
                break

        if scheduler:
            scheduler.step()
